"""
强化学习训练器模块

该模块实现了强化学习算法的训练逻辑，负责从轨迹数据中学习并更新模型参数。
主要功能包括：
1. 轨迹数据采样和批处理
2. 模型训练和参数更新
3. GPU内存管理和错误处理
4. 训练指标记录和可视化

训练器支持内存安全检查，能够在GPU内存不足时自动调整批次大小。
"""

import torch  # PyTorch深度学习框架
import traceback  # 错误追踪工具
import numpy as np  # 数值计算库
from Utils.tensor_ops import _2tensor  # 张量转换工具
from config import GlobalConfig as cfg  # 全局配置
from Utils.gpu_share import GpuShareUnit  # GPU共享单元
from Common.trajectory_sampler import TrajectoryDataSampler  # 轨迹数据采样器


class Trainer():
    """
    强化学习训练器类

    负责管理强化学习模型的训练过程，包括数据采样、损失计算、
    参数更新和性能监控等功能。支持GPU内存管理和错误恢复。
    """

    def __init__(self, model, train_config):
        """
        初始化训练器

        Args:
            model: 要训练的强化学习模型(如QMIX)
            train_config: 训练配置对象，包含各种超参数
        """
        self.policy = model  # 要训练的策略模型
        self.train_epoch = train_config.train_epoch  # 训练轮次
        self.use_avail_act = train_config.action_filter_enabled  # 是否使用可用动作过滤
        self.max_grad_norm = train_config.max_grad_norm  # 最大梯度裁剪范数
        self.memory_safety_check = train_config.memory_safety_check  # 内存安全检查标志
        self.lr = train_config.learning_rate  # 学习率
        self.gamma = train_config.gamma  # 折扣因子

        # 训练统计信息
        self.train_update_cnt = 0  # 训练更新次数计数器
        self.trivial_dict = {}  # 训练指标字典

        # GPU共享单元，用于多进程GPU资源管理
        self.gpu_share_unit = GpuShareUnit(cfg.device, gpu_party=cfg.gpu_party)

    def train_model_on_trajectories(self, traj_pool, task):
        """
        在轨迹数据上训练模型的主入口函数

        该函数处理GPU内存管理，如果内存不足会自动调整批次大小并重试。
        支持内存安全检查模式，能够在OOM(Out of Memory)错误时恢复。

        Args:
            traj_pool: 轨迹池，存储训练数据
            task: 任务标志(通常为'train')
        """
        while True:
            try:
                # 在GPU共享单元的保护下执行训练
                with self.gpu_share_unit:
                    self.learn_from_trajectories(traj_pool, task)
                break  # 成功执行到这说明GPU内存充足，退出循环
            except RuntimeError as err:
                # 捕获运行时错误(通常是GPU内存不足)
                print(traceback.format_exc())

                if self.memory_safety_check:
                    # 内存安全检查模式：调整采样大小并重试
                    # 在某些情况下，单次调整可能不够，需要递归调整
                    if TrajectoryDataSampler.MaxSampleNum[-1] < 0:
                        TrajectoryDataSampler.MaxSampleNum.pop(-1)
                    assert TrajectoryDataSampler.MaxSampleNum[-1] > 0
                    TrajectoryDataSampler.MaxSampleNum[-1] = -1
                    print('Insufficient gpu memory, using previous sample size !')
                else:
                    # 如果未启用内存安全检查，直接抛出断言错误
                    assert False

            # 清空GPU缓存，释放内存
            torch.cuda.empty_cache()

    def learn_from_trajectories(self, traj_pool, task):
        """
        从轨迹数据中学习的核心函数

        执行实际的训练循环，包括数据采样、损失计算和参数更新。

        Args:
            traj_pool: 轨迹池，包含训练数据
            task: 任务标志

        Returns:
            int: 更新后的训练计数
        """
        # 创建轨迹数据采样器
        sampler = TrajectoryDataSampler(n_div=1, traj_pool=traj_pool, flag=task,
                                        memory_safety_check=self.memory_safety_check)

        # 执行多个训练轮次
        for e in range(self.train_epoch):
            # 重置采样器并获取迭代器
            sample_iter = sampler.reset_and_get_iter()
            # 获取下一批训练样本
            train_sample = next(sample_iter)

            # 运行训练步骤，计算损失
            loss_final, loss_dict = self.run_training(task, train_sample)
            # 对损失进行缩放(可能是调优参数)
            loss_final = loss_final * 0.5

            # 第一轮训练时打印GPU内存使用情况
            if e == 0:
                print('[PPO.py] Memory Allocated %.2f GB' % (torch.cuda.memory_allocated() / 1073741824))

            # 记录训练指标
            self.record_training_metrics(dictionary=loss_dict)
            loss_dict = None  # 释放内存

        # 完成所有轮次的更新
        self.print_training_summary()
        self.train_update_cnt += 1
        return self.train_update_cnt

    def record_training_metrics(self, dictionary):
        """
        记录训练指标

        将每轮训练的指标数据存储到字典中，用于后续统计和分析。

        Args:
            dictionary: 包含各种训练指标的字典
        """
        for key in dictionary:
            # 为每个指标创建列表(如果不存在)
            if key not in self.trivial_dict:
                self.trivial_dict[key] = []
            # 提取标量值并添加到列表中
            item = dictionary[key].item() if hasattr(dictionary[key], 'item') else dictionary[key]
            self.trivial_dict[key].append(item)

    def print_training_summary(self, output=True):
        """
        打印训练摘要信息

        计算并显示所有训练指标的平均值，帮助监控训练进度。

        Args:
            output: 是否输出打印信息
        """
        # 将所有指标列表转换为numpy数组
        for key in self.trivial_dict:
            self.trivial_dict[key] = np.array(self.trivial_dict[key])

        # 准备打印缓冲区
        print_buf = ['[Trainer.py] ']

        # 计算每个指标的平均值并添加到打印缓冲区
        for key in self.trivial_dict:
            self.trivial_dict[key] = self.trivial_dict[key].mean()
            print_buf.append(' %s:%.3f, ' % (key, self.trivial_dict[key]))

        # 输出训练摘要
        if output:
            print(''.join(print_buf))

        # 清空指标字典，为下一轮训练做准备
        self.trivial_dict = {}

    def run_training(self, flag, sample):
        """
        执行单次训练步骤

        从样本数据中提取各个组件，调用模型的evaluate_actions方法进行训练。

        Args:
            flag: 训练标志(通常为'train')
            sample: 训练样本字典，包含状态、动作、奖励等信息

        Returns:
            tuple: (损失值, 损失字典)
        """
        # 从样本中提取训练数据并转换为张量
        state = _2tensor(sample['state'])        # 状态张量
        reward = _2tensor(sample['reward'])      # 奖励张量
        action = _2tensor(sample['action'])      # 动作张量
        next_state = _2tensor(sample['next_state'])  # 下一状态张量
        done = _2tensor(sample['done'])          # 完成标志张量
        # 可用动作张量(如果存在)
        avail_act = _2tensor(sample['avail_act']) if 'avail_act' in sample else None

        # 确保是训练模式
        assert flag == 'train'

        # 调用策略模型的evaluate_actions方法进行训练
        # 该方法会计算损失、执行反向传播和参数更新
        loss = self.policy.evaluate_actions(state_batch=state,
                                            action_batch=action,
                                            avail_act_batch=avail_act,
                                            next_state_batch=next_state,
                                            reward_batch=reward,
                                            done_batch=done,
                                            gamma=self.gamma,
                                            max_grad_norm=self.max_grad_norm)

        # 构建损失字典用于记录
        loss_dict = {"value_l1loss": loss}
        return loss, loss_dict